import numpy as np
import matplotlib.pyplot as plt
import os
from scipy.integrate import simps

bin_size= 0.2
min_edge = 0.
max_edge = 45.
bin_list = np.arange(min_edge, max_edge, bin_size)

a = np.loadtxt('dist_type1.txt')
b = np.loadtxt('dist_type2.txt')

fig, ax1 = plt.subplots()
left,bottom,width,height = [0.4,0.4,0.45,0.45]
ax2 = fig.add_axes([left,bottom,width,height])

ax1.plot(a[:,0],a[:,1]/1e6)
ax1.plot(b[:,0],b[:,1]/1e6)
#print(simps(a[:,1]/1e6,a[:,0]),simps(b[:,1]/1e6,b[:,0]))
ax1.set_xlim(0,30)
ax1.set_ylim(0,1.5)
ax1.set_xticks([0,10,20,30])
ax1.set_yticks([0.0,0.5,1.0,1.5])
ax1.minorticks_on()
ax1.set_xlabel(r'$\theta$ (degree)',fontsize=20)
ax1.set_ylabel('Rotation angle distribution (a.u.)',fontsize=20)
ax1.legend(['Type1','Type2'],fontsize=20,frameon=False,loc='lower right',ncol=2,columnspacing=0.05,handlelength=1)
ax1.tick_params(which='both',direction='in',labelsize=20)

a = np.loadtxt('fft_ave_sum1.txt')
b = np.loadtxt('fft_ave_sum2.txt')

n = 132000
timestep = 0.01 # in ps

R = 100
x = a[:,0].reshape(-1, R).mean(axis=1)
a_bin = a[:,1].reshape(-1, R).mean(axis=1)
b_bin = b[:,1].reshape(-1, R).mean(axis=1)

ax2.plot(x[1:],a_bin[1:]*5)
ax2.plot(x[1:],b_bin[1:])
ax2.set_xlim(0,80)
ax2.set_ylim(0,0.06)
ax2.set_xticks([0,20,40,60,80])
ax2.set_yticks([0.00,0.03,0.06])
ax2.minorticks_on()
ax2.set_xlabel('Energy (meV)',fontsize=18)
ax2.set_ylabel('Power spectrum',fontsize=18)
ax2.tick_params(which='both',direction='in',labelsize=18)


plt.savefig('AngleDistributionFFT.pdf',format='pdf',dpi=300,bbox_inches='tight')
plt.show()
